from typing import Dict, Any, Optional, List, Tuple
from pathlib import Path
import json
from datetime import datetime
import time

from config import config
from utils import HTTPClient, LogManager, file_manager

class CloudClient:
    """云端服务客户端"""
    
    def __init__(self):
        self.client = HTTPClient(config.cloud_server.base_url, config.api.timeout)
        self.logger = LogManager.get_logger("CloudClient")
        self.node_id = config.edge_node.id
    
    def register_node(self) -> bool:
        """向云端注册边侧节点"""
        try:
            # 构建节点信息
            node_data = {
                "node_id": self.node_id,
                "name": config.edge_node.name,
                "node_type": "edge",
                "location": config.edge_node.location,
                "status": "online",
                "service_ip": f"{config.edge_node.service_ip}:{config.edge_node.service_port}",
                "cpu": config.edge_node.cpu,
                "memory": config.edge_node.memory,
                "gpu": config.edge_node.gpu,
                "bandwidth": config.edge_node.bandwidth,
                "running_services": [config.edge_node.service_type],
                "deployed_models": [],
                "cpu_load": 0.0,
                "memory_load": 0.0,
                "network_load": 0.0,
                "collected_sample_count": 0,
                "uncollected_sample_count": 0,
                "registration_time": datetime.now().isoformat(),
                "edge_version": "1.0.0"
            }
            
            # 首先检查节点是否已存在
            success, existing_node = self.client.get(f"nodes/{self.node_id}")
            
            if success and existing_node:
                # 节点已存在，更新信息
                self.logger.info(f"节点 {self.node_id} 已存在，更新节点信息")
                update_data = {
                    "status": "online",
                    "service_ip": f"{config.edge_node.service_ip}:{config.edge_node.service_port}",
                    "running_services": [config.edge_node.service_type],
                    "cpu_load": 0.0,
                    "memory_load": 0.0,
                    "network_load": 0.0,
                    "collected_sample_count": 0,
                    "uncollected_sample_count": 0,
                    "last_seen": datetime.now().isoformat()
                }
                
                success, result = self.client.put(f"nodes/{self.node_id}", json_data=update_data)
                if success:
                    self.logger.info(f"节点 {self.node_id} 信息更新成功")
                    return True
                else:
                    self.logger.error(f"更新现有节点失败: {result}")
                    return False
            else:
                # 节点不存在，创建新节点
                success, result = self.client.post("nodes", json_data=node_data)
                
                if success:
                    self.logger.info(f"节点 {self.node_id} 注册成功")
                    return True
                else:
                    self.logger.error(f"节点注册失败: {result}")
                    return False
                
        except Exception as e:
            self.logger.error(f"注册节点异常: {str(e)}")
            return False
    
    def update_node_status(self, status_data: Dict[str, Any]) -> bool:
        """更新节点状态到云端"""
        try:
            success, result = self.client.put(f"nodes/{self.node_id}", json_data=status_data)
            if success:
                self.logger.debug(f"节点状态更新成功: {self.node_id}")
                return True
            else:
                self.logger.debug(f"节点状态更新失败: {result}")
                return False
        except Exception as e:
            self.logger.debug(f"更新节点状态异常: {str(e)}")
            return False
    
    def get_sample_info(self) -> Optional[Dict[str, Any]]:
        """获取样本信息"""
        try:
            dataset_dir = Path(config.dataset.datasets_dir)
            temp_dir = Path(config.dataset.temp_dir)
            confidence_dir = dataset_dir / config.dataset.confidence_dir
            
            # 统计样本信息
            confidence_files = file_manager.list_files(confidence_dir, "*.txt")
            inferred_count = len(confidence_files)
            
            temp_images = file_manager.list_files(temp_dir, "*.jpg") + file_manager.list_files(temp_dir, "*.png")
            uninferred_count = len(temp_images)
            
            # 已归集样本统计
            collected_marker_file = dataset_dir / ".collected_samples"
            collected_samples = []
            if collected_marker_file.exists():
                try:
                    with open(collected_marker_file, 'r') as f:
                        collected_samples = json.load(f)
                except:
                    collected_samples = []
            
            collected_count = len(collected_samples)
            uncollected_count = max(0, inferred_count - collected_count)
            
            # 按难度阈值统计样本数量
            difficulty_stats = self._analyze_confidence_distribution(confidence_files)
            
            return {
                "node_id": self.node_id,
                "task_name": "设备缺陷识别",
                "collected_sample_count": collected_count,
                "uncollected_sample_count": uncollected_count,
                "uninferred_sample_count": uninferred_count,
                "difficulty_threshold_stats": difficulty_stats,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"获取样本信息失败: {str(e)}")
            return None
    
    def upload_samples(self, zip_file_path: Path, max_retries: int = 3) -> Tuple[bool, Dict]:
        """上传样本数据到云端（带指数退避重试）"""
        retry_count = 0
        base_delay = 2  # 基础延迟时间（秒）
        
        while retry_count < max_retries:
            try:
                if not zip_file_path.exists():
                    return False, {"error": "文件不存在"}
                
                with open(zip_file_path, 'rb') as f:
                    files = {'file': (zip_file_path.name, f, 'application/zip')}
                    success, result = self.client.post(
                        "samples/transfer",
                        files=files,
                        params={"source_path": str(zip_file_path), "file_format": "zip"}
                    )
                
                if success:
                    self.logger.info(f"样本上传成功: {zip_file_path}")
                    return True, result
                else:
                    self.logger.error(f"样本上传失败: {result}")
                    if retry_count < max_retries - 1:
                        delay = base_delay * (2 ** retry_count)  # 指数退避
                        self.logger.info(f"准备重试上传，等待 {delay} 秒，剩余次数: {max_retries - retry_count - 1}")
                        time.sleep(delay)
                        retry_count += 1
                        continue
                    return False, result
                    
            except Exception as e:
                self.logger.error(f"上传样本异常: {str(e)}")
                if retry_count < max_retries - 1:
                    delay = base_delay * (2 ** retry_count)  # 指数退避
                    self.logger.info(f"准备重试上传，等待 {delay} 秒，剩余次数: {max_retries - retry_count - 1}")
                    time.sleep(delay)
                    retry_count += 1
                    continue
                return False, {"error": str(e)}
        
        return False, {"error": "达到最大重试次数"}
    
    def upload_samples_with_confirmation(self, zip_file_path: Path, sample_names: List[str], max_retries: int = 3) -> Tuple[bool, Dict]:
        """上传样本数据并等待云端确认"""
        # 首先上传文件
        success, result = self.upload_samples(zip_file_path, max_retries)
        
        if not success:
            return success, result
        
        # 等待云端确认处理完成
        if "task_id" in result or "upload_id" in result:
            task_id = result.get("task_id") or result.get("upload_id")
            return self._wait_for_upload_confirmation(task_id, sample_names)
        
        # 如果没有任务ID，直接返回上传结果
        return success, result
    
    def _wait_for_upload_confirmation(self, task_id: str, sample_names: List[str], timeout: int = 300) -> Tuple[bool, Dict]:
        """等待云端上传确认"""
        start_time = time.time()
        check_interval = 5  # 每5秒检查一次
        
        while time.time() - start_time < timeout:
            try:
                # 检查任务状态
                success, status_result = self.client.get(f"tasks/{task_id}/status")
                
                if success and isinstance(status_result, dict):
                    status = status_result.get("status", "").lower()
                    
                    if status == "completed":
                        self.logger.info(f"云端确认上传成功: {len(sample_names)} 个样本")
                        return True, {
                            "message": "上传已确认",
                            "task_id": task_id,
                            "samples": sample_names
                        }
                    elif status in ["failed", "error"]:
                        error_msg = status_result.get("message", "上传处理失败")
                        self.logger.error(f"云端上传处理失败: {error_msg}")
                        return False, {"error": error_msg}
                    # status为pending或processing时继续等待
                
                time.sleep(check_interval)
                
            except Exception as e:
                self.logger.warning(f"检查上传状态异常: {str(e)}")
                time.sleep(check_interval)
        
        # 超时情况
        self.logger.warning(f"等待云端确认超时，任务ID: {task_id}")
        return False, {"error": f"等待云端确认超时: {timeout}秒"}
    
    def _analyze_confidence_distribution(self, confidence_files: List[Path]) -> Dict[str, int]:
        """分析置信度分布"""
        # 使用sample_manager的统一方法
        from sample_manager import sample_manager
        return sample_manager._analyze_sample_difficulty(confidence_files)
    
    def get_available_models(self) -> Optional[List[Dict]]:
        """获取云端可用的模型列表"""
        try:
            success, result = self.client.get("models")
            if success and result:
                models = result.get("data", [])  # 修正字段名：从"models"改为"data"
                if models:
                    self.logger.info(f"获取到 {len(models)} 个可用模型")
                    for i, model in enumerate(models, 1):
                        model_name = model.get("model_name", "未命名")
                        model_id = model.get("model_id", "N/A")
                        accuracy = model.get("accuracy", model.get("map_50", 0))
                        # 处理accuracy为None的情况
                        accuracy_str = f"{accuracy:.3f}" if accuracy is not None else "N/A"
                        self.logger.info(f"  {i}. [{model_id}] {model_name} (准确率: {accuracy_str})")
                else:
                    self.logger.warning("⚠️ 云端当前没有可用模型!")
                    self.logger.info("请先在云端创建训练任务或上传模型文件")
                return models
            else:
                self.logger.warning(f"获取模型列表失败: {result}")
                return None
        except Exception as e:
            self.logger.error(f"获取模型列表异常: {str(e)}")
            return None
    
    def download_model_from_cloud(self, model_id: str, target_path: Path = None) -> Tuple[bool, str]:
        """从云端下载指定模型"""
        try:
            self.logger.info(f"开始从云端下载模型: {model_id}")
            
            # 确定保存路径
            if target_path is None:
                from pathlib import Path
                models_dir = Path("models")
                models_dir.mkdir(parents=True, exist_ok=True)
                target_path = models_dir / f"{model_id}.pt"
            
            # 直接使用requests进行流式下载，避免HTTPClient的限制
            try:
                import requests
                
                # 构建完整的下载URL - 修正URL重复问题
                # config.cloud_server.base_url 已经包含 /api/v1，所以不需要再次添加
                full_url = f"{config.cloud_server.base_url}/models/{model_id}/download"
                
                self.logger.debug(f"下载URL: {full_url}")
                
                response = requests.get(full_url, stream=True, timeout=300)
                response.raise_for_status()
                
                # 下载到临时文件，然后原子性移动
                import tempfile
                with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
                    for chunk in response.iter_content(chunk_size=8192):
                        tmp_file.write(chunk)
                    temp_path = Path(tmp_file.name)
                
                # 原子性移动文件
                import shutil
                shutil.move(str(temp_path), str(target_path))
                
                self.logger.info(f"模型下载成功: {target_path}")
                return True, str(target_path)
                
            except Exception as e:
                self.logger.error(f"下载模型文件失败: {str(e)}")
                return False, f"下载文件失败: {str(e)}"
                
        except Exception as e:
            self.logger.error(f"下载模型异常: {str(e)}")
            return False, f"下载异常: {str(e)}"
    
    def get_model_by_service_type(self, service_type: str) -> Tuple[bool, Optional[str]]:
        """根据服务类型从云端获取相应的模型"""
        try:
            self.logger.info(f"根据服务类型获取模型: {service_type}")
            
            # 获取可用模型列表
            models = self.get_available_models()
            if not models:
                return False, "云端没有可用模型，请先在云端创建训练任务或上传模型"
            
            if len(models) == 0:
                self.logger.warning("🔍 云端模型列表为空，无法获取模型")
                self.logger.info("💡 解决方案:")
                self.logger.info("   1. 在云端启动训练任务生成新模型")
                self.logger.info("   2. 手动上传预训练模型到云端")
                self.logger.info("   3. 检查云端数据库中的模型记录")
                return False, "云端没有可用模型"
            
            # 根据服务类型筛选模型
            suitable_models = []
            for model in models:
                model_tags = model.get("tags", [])
                model_description = model.get("description", "")
                
                # 根据服务类型匹配模型
                if self._is_model_suitable_for_service(model, service_type):
                    suitable_models.append(model)
            
            if not suitable_models:
                # 如果没有专门匹配的模型，使用第一个可用模型作为通用模型
                self.logger.warning(f"没有找到专门适用于 '{service_type}' 的模型，使用通用模型")
                selected_model = models[0]
                self.logger.info(f"🎯 选择通用模型: {selected_model.get('model_id', 'N/A')}")
            else:
                # 选择最适合的模型（优先选择最新的）
                selected_model = max(suitable_models, key=lambda m: m.get("updated_at", ""))
                self.logger.info(f"🎯 找到适用于 '{service_type}' 的模型: {selected_model.get('model_id', 'N/A')}")
            
            model_id = selected_model.get("model_id")
            if not model_id:
                return False, "模型列表中没有有效的模型ID"
            
            # 下载模型
            success, result = self.download_model_from_cloud(model_id)
            
            if success:
                self.logger.info(f"✅ 成功获取适用于 '{service_type}' 的模型: {result}")
                return True, result
            else:
                return False, result
                
        except Exception as e:
            self.logger.error(f"根据服务类型获取模型失败: {str(e)}")
            return False, f"获取模型异常: {str(e)}"
    
    def _is_model_suitable_for_service(self, model: Dict, service_type: str) -> bool:
        """判断模型是否适用于指定的服务类型"""
        model_tags = model.get("tags", [])
        model_description = model.get("description", "").lower()
        model_name = model.get("model_name", "").lower()

      # self.logger.info(f"service_type:{service_type} , model_tags: {model_tags}")
        if service_type in model_tags:
           return True
        return False
    
    def get_default_model_from_cloud(self) -> Tuple[bool, Optional[str]]:
        """从云端获取并下载默认推荐模型"""
        try:
            # 获取可用模型列表
            models = self.get_available_models()
            if not models:
                return False, "云端没有可用模型，请先在云端创建训练任务或上传模型"
            
            if len(models) == 0:
                self.logger.warning("🔍 云端模型列表为空，无法下载默认模型")
                self.logger.info("💡 解决方案:")
                self.logger.info("   1. 在云端启动训练任务生成新模型")
                self.logger.info("   2. 手动上传预训练模型到云端")
                self.logger.info("   3. 检查云端数据库中的模型记录")
                return False, "云端没有可用模型"
            
            # 选择第一个可用模型作为默认模型
            default_model = models[0]
            model_id = default_model.get("model_id")
            
            if not model_id:
                return False, "模型列表中没有有效的模型ID"
            
            self.logger.info(f"🎯 选择默认模型进行下载: {model_id}")
            
            # 下载模型
            success, result = self.download_model_from_cloud(model_id)
            
            if success:
                return True, result
            else:
                return False, result
                
        except Exception as e:
            self.logger.error(f"获取默认模型失败: {str(e)}")
            return False, f"获取默认模型异常: {str(e)}"

    def check_cloud_connection(self) -> bool:
        """检查与云端的连接状态"""
        try:
            # 尝试访问健康检查接口
            success, result = self.client.get("health")
            if success:
                self.logger.debug("云端连接正常")
                return True
            else:
                self.logger.debug(f"云端连接检查失败: {result}")
                return False
        except Exception as e:
            self.logger.debug(f"云端连接检查异常: {str(e)}")
            return False

# 全局实例
cloud_client = CloudClient() 
